﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Runtime.InteropServices;
using System.Collections.ObjectModel;
using System.IO;

namespace Luminal
{
    public enum BuildStatus
    {
        Success = 0,
        None = -1,
        Error = -2,
        InProgress = -3,
    }

    public struct ProgramBinary
    {
        public DeviceIdentifier Device;
        public byte[] BinaryData;

        public ProgramBinary(DeviceIdentifier device, byte[] data)
        {
            Device = device;
            BinaryData = data;
        }

        public static ProgramBinary FromFile(string filename, DeviceIdentifier device)
        {
            return new ProgramBinary(device, File.ReadAllBytes(filename));
        }
    }

    public sealed class Program : LuminalObject
    {
        #region Interop
        // clProgram
        [DllImport("OpenCL.dll")]
        static extern IntPtr clCreateProgramWithSource(IntPtr context,
            uint count,
            string[] strings,
            SizeT[] lengths,
            out ErrorCode errcode_ret);

        [DllImport("OpenCL.dll")]
        static extern IntPtr clCreateProgramWithBinary(IntPtr context,
            uint num_devices,
            IntPtr[] device_list,
            SizeT[] lengths,
            IntPtr[] binaries,
            ErrorCode[] binary_status,
            out ErrorCode errcode_ret);

        [DllImport("OpenCL.dll")]
        static extern ErrorCode clBuildProgram(IntPtr program,
            uint num_devices,
            IntPtr[] device_list,
            string options,
            IntPtr pfnNotify,
            IntPtr userData); 

        [DllImport("OpenCL.dll")]
        static extern ErrorCode clGetProgramInfo(IntPtr program,
            cl_program_info param_name,
            SizeT param_value_size,
            IntPtr param_value,
            out SizeT param_value_size_ret);

        [DllImport("OpenCL.dll")]
        static extern ErrorCode clGetProgramBuildInfo(IntPtr program,
            IntPtr device,
            cl_program_build_info param_name,
            SizeT param_value_size,
            IntPtr param_value,
            out SizeT param_value_size_ret);

        [DllImport("OpenCL.dll")]
        static extern ErrorCode clReleaseProgram(IntPtr program);

        [DllImport("OpenCL.dll")]
        static extern ErrorCode clUnloadCompiler();

        // clKernel
        [DllImport("OpenCL.dll")]
        static extern ErrorCode clCreateKernelsInProgram(IntPtr program,
            uint num_kernels,
            IntPtr[] kernels,
            out uint num_kernels_ret);
        #endregion

        #region Properties
        public Context Context { get; private set; }

        public ReadOnlyCollection<DeviceIdentifier> Devices
        {
            get
            {
                IntPtr[] devicePtrs = GetInfo<IntPtr>(cl_program_info.CL_PROGRAM_DEVICES, (int)GetInfo<uint>(cl_program_info.CL_PROGRAM_NUM_DEVICES));
                List<DeviceIdentifier> devices = new List<DeviceIdentifier>();
                foreach (IntPtr dPtr in devicePtrs)
                    devices.Add(DeviceIdentifier.FromPointer(dPtr));
                return devices.AsReadOnly();
            }
        }

        KernelCollection kernels;
        public KernelCollection Kernels
        {
            get
            {
                if (kernels == null)
                {
                    uint kernelCount;
                    Util.ThrowError(
                        "Could not create kernels in program.",
                        clCreateKernelsInProgram(InternalPointer, 0, null, out kernelCount));

                    IntPtr[] outputPtrs = new IntPtr[kernelCount];
                    Util.ThrowError(
                        "Could not create kernels in program.",
                        clCreateKernelsInProgram(InternalPointer, kernelCount, outputPtrs, out kernelCount));

                    kernels = new KernelCollection(from ptr in outputPtrs
                                                   select Kernel.FromPointer(ptr));
                }

                return kernels;
            }
        }

        private T GetInfo<T>(cl_program_info info)
        {
            Type type = typeof(T);
            if (type.IsEnum)
                type = Enum.GetUnderlyingType(type);

            SizeT dataSize;
            Util.ThrowError(
                "Could not get program info.",
                clGetProgramInfo(InternalPointer, info, 0, IntPtr.Zero, out dataSize));

            if (type.IsValueType)
            {
                byte[] outputBytes = new byte[dataSize];
                Util.Pin(outputBytes, ptr =>
                {
                    Util.ThrowError(
                        "Could not get program info.",
                        clGetProgramInfo(InternalPointer, info, dataSize, ptr, out dataSize));
                });

                IntPtr outputPtr = Marshal.AllocHGlobal(outputBytes.Length);
                Marshal.Copy(outputBytes, 0, outputPtr, outputBytes.Length);

                T output = (T)Marshal.PtrToStructure(outputPtr, type);
                Marshal.FreeHGlobal(outputPtr);

                return output;
            }
            else if (type == typeof(string))
            {
                IntPtr outputPtr = Marshal.AllocHGlobal((int)dataSize);
                Util.ThrowError(
                    "Could not get program info.",
                    clGetProgramInfo(InternalPointer, info, dataSize, outputPtr, out dataSize));

                string output = Marshal.PtrToStringAnsi(outputPtr);
                Marshal.FreeHGlobal(outputPtr);

                return (T)(object)output;
            }
            else
            {
                throw new Exception("Type not supported.");
            }
        }

        private T GetBuildInfo<T>(cl_program_build_info info, DeviceIdentifier device)
        {
            Type type = typeof(T);
            if (type.IsEnum)
                type = Enum.GetUnderlyingType(type);

            SizeT dataSize;
            Util.ThrowError(
                "Could not get program info.",
                clGetProgramBuildInfo(InternalPointer, device.InternalPointer, info, 0, IntPtr.Zero, out dataSize));

            if (type.IsValueType)
            {
                byte[] outputBytes = new byte[dataSize];
                Util.Pin(outputBytes, ptr =>
                {
                    Util.ThrowError(
                        "Could not get program info.",
                        clGetProgramBuildInfo(InternalPointer, device.InternalPointer, info, dataSize, ptr, out dataSize));
                });

                IntPtr outputPtr = Marshal.AllocHGlobal(outputBytes.Length);
                Marshal.Copy(outputBytes, 0, outputPtr, outputBytes.Length);

                T output = (T)Marshal.PtrToStructure(outputPtr, type);
                Marshal.FreeHGlobal(outputPtr);

                return output;
            }
            else if (type == typeof(string))
            {
                IntPtr outputPtr = Marshal.AllocHGlobal((int)dataSize);
                Util.ThrowError(
                    "Could not get program info.",
                    clGetProgramBuildInfo(InternalPointer, device.InternalPointer, info, dataSize, outputPtr, out dataSize));

                string output = Marshal.PtrToStringAnsi(outputPtr);
                Marshal.FreeHGlobal(outputPtr);

                return (T)(object)output;
            }
            else
            {
                throw new Exception("Type not supported.");
            }
        }

        private T[] GetInfo<T>(cl_program_info info, int count) where T : struct
        {
            Type type = typeof(T);
            if (type.IsEnum)
                type = Enum.GetUnderlyingType(type);

            SizeT dataSize;
            T[] output = new T[count];
            Util.Pin(output, ptr =>
                {
                    Util.ThrowError(
                        "Could not get program info.",
                        clGetProgramInfo(InternalPointer, info, count * Marshal.SizeOf(typeof(T)), ptr, out dataSize));
                });

            return output;
        }
        #endregion

        #region Constructor
        private Program(IntPtr ptr)
        {
            Initialize(ptr);

            try
            {
                Context = Context.FromPointer(GetInfo<IntPtr>(cl_program_info.CL_PROGRAM_CONTEXT));
            }
            catch (Exception)
            {
                Dispose();
                throw;
            }
        }
        #endregion

        #region Static Methods
        public static Program FromSource(Context context, params string[] source)
        {
            ErrorCode errorCode;
            IntPtr ptr = clCreateProgramWithSource(context.InternalPointer, (uint)source.Length, source, null, out errorCode);

            Util.ThrowError("Could not create program from source.", errorCode);

            return new Program(ptr);
        }

        public static Program FromFile(Context context, string filename)
        {
            return FromSource(context, File.ReadAllText(filename));
        }

        public static Program FromBinaries(Context context, params ProgramBinary[] binaries)
        {
            IntPtr[] binaryData = new IntPtr[binaries.Length];
            GCHandle[] binaryHandles = new GCHandle[binaries.Length];
            SizeT[] lengths = new SizeT[binaries.Length];
            IntPtr[] devicePtrs = new IntPtr[binaries.Length];

            for (int i = 0; i < binaries.Length; i++)
            {
                binaryHandles[i] = GCHandle.Alloc(binaries[i].BinaryData, GCHandleType.Pinned);
                binaryData[i] = binaryHandles[i].AddrOfPinnedObject();
                lengths[i] = binaries[i].BinaryData.Length;
                devicePtrs[i] = binaries[i].Device.InternalPointer;
            }

            ErrorCode errorCode;
            ErrorCode[] binaryStatuses = new ErrorCode[binaries.Length];
            IntPtr ptr = clCreateProgramWithBinary(context.InternalPointer, (uint)devicePtrs.Length, devicePtrs, lengths, binaryData, binaryStatuses, out errorCode);

            Util.ThrowError("Could not create program from binaries.", errorCode);

            return new Program(ptr);
        }

        public static void UnloadCompiler()
        {
            Util.ThrowError(
                "Could not unload compiler.",
                clUnloadCompiler());
        }
        #endregion

        #region Methods
        public bool Build(IEnumerable<DeviceIdentifier> devices, string options, out string[] buildLogs)
        {
            IntPtr[] devicePtrs = Util.ExtractPointers<DeviceIdentifier>(devices);
            ErrorCode errorCode = clBuildProgram(InternalPointer, (uint)devicePtrs.Length, devicePtrs, options, IntPtr.Zero, IntPtr.Zero);

            buildLogs = new string[devicePtrs.Length];
            for (int i = 0; i < buildLogs.Length; i++)
                buildLogs[i] = GetBuildInfo<string>(cl_program_build_info.CL_PROGRAM_BUILD_LOG, devices.ElementAt(i));

            if (errorCode != ErrorCode.CL_BUILD_PROGRAM_FAILURE)
                Util.ThrowError("Could not build program.", errorCode);
            else
                return false;

            return true;
        }

        public void Build(IEnumerable<DeviceIdentifier> devices, string options)
        {
            IntPtr[] devicePtrs = Util.ExtractPointers<DeviceIdentifier>(devices);
            ErrorCode errorCode = clBuildProgram(InternalPointer, (uint)devicePtrs.Length, devicePtrs, options, IntPtr.Zero, IntPtr.Zero);

            Util.ThrowError("Could not build program.", errorCode);
        }

        public BuildStatus GetBuildStatus(DeviceIdentifier device)
        {
            return GetBuildInfo<BuildStatus>(cl_program_build_info.CL_PROGRAM_BUILD_STATUS, device);
        }

        public string GetSource()
        {
            return GetInfo<string>(cl_program_info.CL_PROGRAM_SOURCE);
        }

        public ProgramBinary[] GetBinaries()
        {
            SizeT[] binarySizes = GetInfo<SizeT>(cl_program_info.CL_PROGRAM_BINARY_SIZES, Devices.Count);

            IntPtr[] binaryPtrs = new IntPtr[binarySizes.Length];
            for (int i = 0; i < binaryPtrs.Length; i++)
                binaryPtrs[i] = Marshal.AllocHGlobal((int)binarySizes[i]);

            SizeT garbage;

            Util.Pin(binaryPtrs, ptr =>
                {
                    Util.ThrowError(
                        "Could not get program binaries.",
                        clGetProgramInfo(InternalPointer, cl_program_info.CL_PROGRAM_BINARIES, Marshal.SizeOf(typeof(IntPtr)), ptr, out garbage));
                });

            ProgramBinary[] output = new ProgramBinary[binaryPtrs.Length];
            for (int i = 0; i < output.Length; i++)
            {
                byte[] data = new byte[binarySizes[i]];
                Marshal.Copy(binaryPtrs[i], data, 0, data.Length);
                output[i] = new ProgramBinary(Devices[i], data);
            }
            return output;
        }
        #endregion

        protected override void Dispose(bool disposing)
        {
            if (disposing)
            {
                if (kernels != null)
                    foreach (Kernel kernel in kernels)
                        if (!kernel.Disposed)
                            kernel.Dispose();

                Util.ThrowError("Could not release program.", clReleaseProgram(InternalPointer));
            }

            base.Dispose(disposing);
        }

        public static Program FromPointer(IntPtr ptr)
        {
            return ObjectTracker.ContainsObject(ptr) ? ObjectTracker.GetObject<Program>(ptr) : new Program(ptr);
        }
    }
}

